# YOLOv8检测核心类
from ultralytics import YOLO
import numpy as np
from typing import List, Dict, Any, Tuple
import logging
import torch


class YOLODetector:
    """YOLO目标检测器"""

    def __init__(self, model_config: Dict[str, Any]):
        """
        初始化检测器
        Args:
            model_config: 模型配置字典
        """
        self.config = model_config['model']
        self.performance_config = model_config['performance']

        # 加载模型
        self.model = YOLO(self.config['path'])

        # COCO数据集类别名称
        self.class_names = {
            0: 'person',
            39: 'laptop',  # 电脑
            41: 'cup'  # 杯子
        }

        logging.info(f"YOLODetector初始化完成，模型路径: {self.config['path']}")
        print(f"PyTorch 版本: {torch.__version__}")
        print(f"CUDA 可用: {torch.cuda.is_available()}")
        print(f"GPU 名称: {torch.cuda.get_device_name(0)}")
        print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
        print(f"Current device: {torch.cuda.current_device()}")
        print(f"Device name: {torch.cuda.get_device_name(0)}")

    def detect(self, frame: np.ndarray) -> List[Dict[str, Any]]:
        """
        执行目标检测
        Args:
            frame: 输入图像帧
        Returns:
            检测结果列表
        """
        try:
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            # 执行推理
            results = self.model(
                frame,
                imgsz=self.config['imgsz'],  # 输入尺寸 (如640)
                conf=0.4,  # 置信度阈值
                # device=self.config['device'],   # 强制使用第一块GPU
                device=device,   # 强制使用第一块GPU
                classes=self.performance_config['classes'], # 可选：指定检测类别
                max_det=self.performance_config['max_det'], # 最大检测数
                half=self.config['half']    # 启用FP16加速（需GPU支持）
            )

            detections = []

            # 解析检测结果
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        confidence = float(box.conf[0].cpu().numpy())
                        class_id = int(box.cls[0].cpu().numpy())

                        # 只处理我们关心的类别
                        if class_id in self.class_names:
                            detection = {
                                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                                'confidence': confidence,
                                'class_id': class_id,
                                'class_name': self.class_names[class_id]
                            }
                            detections.append(detection)

            print("当前推理设备：", device)
            return detections

        except Exception as e:
            logging.error(f"检测过程中出错: {e}")
            return []

    def filter_detections(self, detections: List[Dict[str, Any]],
                          confidence_threshold: float = 0.5) -> List[Dict[str, Any]]:
        """
        过滤检测结果
        Args:
            detections: 原始检测结果
            confidence_threshold: 置信度阈值
        Returns:
            过滤后的检测结果
        """
        return [det for det in detections if det['confidence'] >= confidence_threshold]
